package core

type LessFunc[T any] func(a, b T) bool

type PriorityQueue[T any] struct {
	heap []T
	size int
	less LessFunc[T]
}

func NewPriorityQueue[T any](less LessFunc[T]) *PriorityQueue[T] {
	return &PriorityQueue[T]{
		heap: make([]T, 0),
		size: 0,
		less: less,
	}
}

func (pq *PriorityQueue[T]) Size() int {
	return pq.size
}

func (pq *PriorityQueue[T]) IsEmpty() bool {
	return pq.size == 0
}

func (pq *PriorityQueue[T]) swap(i, j int) {
	pq.heap[i], pq.heap[j] = pq.heap[j], pq.heap[i]
}

// 上浮
func (pq *PriorityQueue[T]) heapifyUp(index int) {
	for index > 0 {
		parentIndex := (index - 1) / 2
		if !pq.less(pq.heap[index], pq.heap[parentIndex]) {
			break
		}
		pq.swap(index, parentIndex)
		index = parentIndex
	}
}

func (pq *PriorityQueue[T]) heapifyDown(index int) {
	for {
		leftChildIndex := index*2 + 1
		rightChildIndex := index*2 + 2
		targetIndex := index

		if leftChildIndex < pq.size && pq.less(pq.heap[leftChildIndex], pq.heap[targetIndex]) {
			targetIndex = leftChildIndex
		}
		if rightChildIndex < pq.size && pq.less(pq.heap[rightChildIndex], pq.heap[targetIndex]) {
			targetIndex = rightChildIndex
		}
		if targetIndex == index {
			break
		}

		pq.swap(index, targetIndex)
		index = targetIndex
	}
}

func (pq *PriorityQueue[T]) Push(value T) {
	pq.heap = append(pq.heap, value)
	pq.size++
	pq.heapifyUp(pq.size - 1)
}

func (pq *PriorityQueue[T]) Pop() (T, bool) {
	var zero T
	if pq.size == 0 {
		return zero, false
	}
	top := pq.heap[0]
	pq.heap[0] = pq.heap[pq.size-1]
	pq.size--
	pq.heap = pq.heap[:pq.size]
	if pq.size > 0 {
		pq.heapifyDown(0)
	}
	return top, true
}

func (pq *PriorityQueue[T]) Peek() (T, bool) {
	var zero T
	if pq.size == 0 {
		return zero, false
	}
	return pq.heap[0], true
}

// 获取全部元素（只读副本）
func (pq *PriorityQueue[T]) All() []T {
	out := make([]T, pq.size)
	copy(out, pq.heap)
	return out
}

// 按条件删除元素
func (pq *PriorityQueue[T]) RemoveIf(match func(T) bool) int {
	n := 0
	i := 0
	for i < pq.size {
		if match(pq.heap[i]) {
			pq.heap[i] = pq.heap[pq.size-1]
			pq.size--
			pq.heap = pq.heap[:pq.size]
			if pq.size > 0 {
				pq.heapifyDown(i)
			}
			n++
		} else {
			i++
		}
	}
	return n
}
